import os
import pandas as pd
from tqdm import tqdm
import argparse
from concurrent.futures import ProcessPoolExecutor, as_completed
from .parser import do_parse


def get_files(work_dir):
    file_hash_set = set()
    for root, _, files in os.walk(work_dir):
        for filename in files:
            file_path = os.path.join(root, filename)
            if filename.endswith(('.pdf', '.PDF', '.png', '.jpg', '.bmp', '.jpeg')):
                with open(file_path, 'rb') as f:
                    file_hash = hash(f.read())
                if file_path not in file_hash_set:
                    file_hash_set.add(file_hash)
                    yield file_path
                

def tasks(work_dir, output):
    with ProcessPoolExecutor() as pool:
        df = pd.DataFrame(columns=['文件名', '路径', '状态', '解析源', '备注', 'inv_code', 'inv_num', 'inv_date', 'inv_chk', 'inv_money'])
        print('提交任务...')
        files = list(get_files(work_dir))
        workers = [pool.submit(do_parse, file_path) for file_path in tqdm(files)]
        print('任务执行中...')
        for t in tqdm(as_completed(workers), total=len(files)):
            state, msg, file_path, data, source = t.result()
            file_name= os.path.basename(file_path)
            row = [file_name, file_path, state, source, msg]
            if not state:
                row += ['', '', '', '', '']
            else:
                row += list(data)
            df = df.append(pd.DataFrame(row, index=df.columns).T)
        df.astype(str).to_excel(output, index=False)
        print('任务执行完成！')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-p', '--path', type=str, required=True, help='输入文件路径')
    parser.add_argument('-o', '--output', type=str, default='发票基础信息.xlsx', help='输入文件路径')
    args = parser.parse_args()
    tasks(args.path, args.output)
